#include "../test_helpers.h"
#include "MessageDispatcher.h"
#include <spdlog/spdlog.h>
#include <gtest/gtest.h>

using namespace MessageFramework;
using namespace MessageFramework::Test;

class MessageDispatcherTest : public ::testing::Test {
protected:
    std::string secretKey = "test-secret-key-32bytes-long!!!";

    void SetUp() override {
        spdlog::set_level(spdlog::level::debug);
    }
};

// ==================== Handler注册测试 ====================

TEST_F(MessageDispatcherTest, RegisterHandler) {
    MessageDispatcher dispatcher;

    auto handler = std::make_shared<MockHandler>("TestMessage");
    dispatcher.registerHandler("TestMessage", handler);

    EXPECT_EQ(dispatcher.getHandlerCount(), 1);

    auto types = dispatcher.getRegisteredTypes();
    EXPECT_EQ(types.size(), 1);
    EXPECT_EQ(types[0], "TestMessage");
}

TEST_F(MessageDispatcherTest, RegisterMultipleHandlers) {
    MessageDispatcher dispatcher;

    auto handler1 = std::make_shared<MockHandler>("Message1");
    auto handler2 = std::make_shared<MockHandler>("Message2");

    dispatcher.registerHandler("Message1", handler1);
    dispatcher.registerHandler("Message2", handler2);

    EXPECT_EQ(dispatcher.getHandlerCount(), 2);
}

TEST_F(MessageDispatcherTest, UnregisterHandler) {
    MessageDispatcher dispatcher;

    auto handler = std::make_shared<MockHandler>("TestMessage");
    dispatcher.registerHandler("TestMessage", handler);
    EXPECT_EQ(dispatcher.getHandlerCount(), 1);

    dispatcher.unregisterHandler("TestMessage");
    EXPECT_EQ(dispatcher.getHandlerCount(), 0);
}

// ==================== 消息分发测试 ====================

TEST_F(MessageDispatcherTest, DispatchToHandler) {
    MessageDispatcher::Config config;
    config.secretKey = secretKey;
    config.verifySignature = true;

    MessageDispatcher dispatcher(config);

    auto handler = std::make_shared<MockHandler>("TestMessage");
    dispatcher.registerHandler("TestMessage", handler);

    // 创建并签名消息
    MessageHeader header = TestHelpers::createTestHeader("TestMessage", "TestPublisher");
    std::string body = TestHelpers::createTestBody();
    SignatureHelper::sign(header, body, secretKey);

    Message message(header, body);

    // 分发消息
    auto result = dispatcher.dispatch(message);

    EXPECT_EQ(result.code, 0);
    EXPECT_EQ(result.handlerName, "MockHandler");
    EXPECT_EQ(handler->getCallCount(), 1);
}

TEST_F(MessageDispatcherTest, DispatchUnknownMessageType) {
    MessageDispatcher dispatcher;

    MessageHeader header = TestHelpers::createTestHeader("UnknownType", "TestPublisher");
    std::string body = TestHelpers::createTestBody();

    Message message(header, body);

    auto result = dispatcher.dispatch(message);

    EXPECT_EQ(result.code, -4);  // No handler found
    EXPECT_FALSE(result.error.empty());
}

TEST_F(MessageDispatcherTest, DispatchWithDefaultHandler) {
    MessageDispatcher dispatcher;

    // 设置默认Handler
    auto defaultHandler = std::make_shared<MockHandler>("DefaultMessage");
    dispatcher.setDefaultHandler(defaultHandler);

    MessageHeader header = TestHelpers::createTestHeader("UnknownType", "TestPublisher");
    std::string body = TestHelpers::createTestBody();

    Message message(header, body);

    auto result = dispatcher.dispatch(message);

    EXPECT_EQ(result.code, 0);
    EXPECT_EQ(defaultHandler->getCallCount(), 1);
}

// ==================== 签名验证测试 ====================

TEST_F(MessageDispatcherTest, DispatchWithSignatureVerification) {
    MessageDispatcher::Config config;
    config.secretKey = secretKey;
    config.verifySignature = true;

    MessageDispatcher dispatcher(config);

    auto handler = std::make_shared<MockHandler>("TestMessage");
    dispatcher.registerHandler("TestMessage", handler);

    MessageHeader header = TestHelpers::createTestHeader("TestMessage", "TestPublisher");
    std::string body = TestHelpers::createTestBody();

    // 不签名，直接分发
    Message message(header, body);
    auto result = dispatcher.dispatch(message);

    // 应该验证失败
    EXPECT_EQ(result.code, -2);  // Signature verification failed
    EXPECT_EQ(handler->getCallCount(), 0);  // Handler不应被调用
}

TEST_F(MessageDispatcherTest, DispatchWithValidSignature) {
    MessageDispatcher::Config config;
    config.secretKey = secretKey;
    config.verifySignature = true;

    MessageDispatcher dispatcher(config);

    auto handler = std::make_shared<MockHandler>("TestMessage");
    dispatcher.registerHandler("TestMessage", handler);

    MessageHeader header = TestHelpers::createTestHeader("TestMessage", "TestPublisher");
    std::string body = TestHelpers::createTestBody();

    // 正确签名
    SignatureHelper::sign(header, body, secretKey);

    Message message(header, body);
    auto result = dispatcher.dispatch(message);

    // 应该成功
    EXPECT_EQ(result.code, 0);
    EXPECT_EQ(handler->getCallCount(), 1);
}

// ==================== 加密消息测试 ====================

TEST_F(MessageDispatcherTest, DispatchEncryptedMessage) {
    std::string encryptionKey = EncryptionHelper::generateKey();

    MessageDispatcher::Config config;
    config.secretKey = secretKey;
    config.encryptionKey = encryptionKey;
    config.verifySignature = true;
    config.autoDecrypt = true;

    MessageDispatcher dispatcher(config);

    auto handler = std::make_shared<MockHandler>("TestMessage");
    dispatcher.registerHandler("TestMessage", handler);

    MessageHeader header = TestHelpers::createTestHeader("TestMessage", "TestPublisher");
    std::string plaintext = "{\"data\":\"secret\"}";

    // 加密消息体
    std::string ciphertext = EncryptionHelper::encrypt(plaintext, encryptionKey);
    header.encrypted = true;
    std::strcpy(header.encryptionAlgo, "AES-256-GCM");

    // 签名（对加密后的body签名）
    SignatureHelper::sign(header, ciphertext, secretKey);

    Message message(header, ciphertext);

    // 分发消息
    auto result = dispatcher.dispatch(message);

    EXPECT_EQ(result.code, 0);
    EXPECT_EQ(handler->getCallCount(), 1);

    // 验证解密后的消息体
    Message lastMsg = handler->getLastMessage();
    EXPECT_EQ(lastMsg.body, plaintext);
    EXPECT_FALSE(lastMsg.header.encrypted);  // 应该被标记为已解密
}

// ==================== 错误处理测试 ====================

TEST_F(MessageDispatcherTest, DispatchInvalidHeader) {
    MessageDispatcher dispatcher;

    auto handler = std::make_shared<MockHandler>("TestMessage");
    dispatcher.registerHandler("TestMessage", handler);

    // 创建无效的Header（空messageId）
    MessageHeader header;
    header.messageId[0] = '\0';
    std::strcpy(header.messageType, "TestMessage");

    Message message(header, "test body");

    auto result = dispatcher.dispatch(message);

    EXPECT_EQ(result.code, -1);  // Invalid message header
    EXPECT_EQ(handler->getCallCount(), 0);
}

int main(int argc, char** argv) {
    ::testing::InitGoogleTest(&argc, argv);
    return RUN_ALL_TESTS();
}
